#!/usr/bin/python3

import argparse
import json
import shutil
import jsonschema
import psutil
import os
import re
import subprocess

DATA_SCHEMA_DIR = "/usr/share/azurelinux-sysinfo"
DATA_SCHEMA_VERSION = "v1"
DATA_SCHEMA_FILENAME = f"sysinfo-schema-{DATA_SCHEMA_VERSION}.json"
LOG_FILE_PATH = "/var/log/azurelinux-sysinfo.log"
SERVICE_NAME = "azurelinux-sysinfo-service"


# This function converts a string that matches
# regex = r"(\d+(\.\d+)?)(min|s|ms)" to seconds
def convert_to_secs(line):
    regex = r"(\d+(?:\.\d+)?)(min|s|ms)"
    time_secs = 0
    for match in re.findall(regex, line):
        time = float(match[0])
        unit = match[1]
        if unit == "min":
            time *= 60
        elif unit == "ms":
            time /= 1000
        time_secs += time
    return time_secs


def collect_os_info():
    print("Collecting os info...")

    release_data = {}
    release_info = subprocess.run(
        ["cat", "/etc/os-release"], capture_output=True, text=True
    )

    kernel_info = subprocess.run(["uname", "-r"], capture_output=True, text=True)

    for line in release_info.stdout.strip().splitlines():
        name, value = line.split("=", maxsplit=1)
        release_data[name] = value.strip('"')


    os_info = {
        "kernel_version": kernel_info.stdout.strip(),
        "release_version": release_data["VERSION"],
        "release_version_id": release_data["VERSION_ID"],
    }

    return os_info


def collect_boot_info():
    print("Collecting boot info...")
    # Known issue: In SELinux enforcing mode, systemd-analyze commands are expected to fail until required policies are added.
    # In this case, the boot times will be 0 and longest running processes will be empty.
    
    # Collect boot time
    result = subprocess.run(["systemd-analyze", "time"], capture_output=True, text=True)

    # Sample output for livecd image:
    # Startup finished in 153ms (firmware) + 554ms (loader) + 1.413s (kernel) + 908ms (userspace) = 3.030s 
    # multi-user.target reached after 897ms in userspace
    # Sample output for host images:
    # Startup finished in 12.688s (kernel) + 8.082s (initrd) + 1min 1.458s (userspace) = 1min 22.230s
    # multi-user.target reached after 1min 966ms in userspace
    
    lines = result.stdout.strip().splitlines()

    # In a test setup on qemu, systemd-analyze returns empty
    if len(lines) < 1 or not(lines[0].startswith("Startup finished in")):
        boot_info = {
            "boot_time": {
                "kernel_boot_time_secs": 0,
                "userspace_boot_time_secs": 0,
                "total_boot_time_secs": 0,
            },
            "longest_running_processes": [],
        }
        return boot_info

    # Define regular expression to extract times 
    timeRegex = r"((?:\d+)(?:\d*min\s?)?(?:\d*\.?\d*s\s?)?(?:\d*\.?\d*ms)?)"
    # Define regular expression to extract values between parentheses
    betweenParenthesesRegex = r"\((.*?)\)"

    boot_time_keys = re.findall(betweenParenthesesRegex, lines[0])
    boot_times = re.findall(timeRegex, lines[0])
    boot_times = [t.strip() for t in boot_times]    
    boot_times_secs = [convert_to_secs(time) for time in boot_times]

    boot_time = dict()
    suffix = "_boot_time_secs"
    for i in range(len(boot_time_keys)):
        bootTimeKey = boot_time_keys[i] + suffix
        bootTimeValue = boot_times_secs[i]
        boot_time[bootTimeKey] = bootTimeValue
    boot_time["total_boot_time_secs"] = boot_times_secs[-1]
    
    # Collect boot time longest running processes
    top_n = 3
    result = subprocess.run(
        ["systemd-analyze", "blame"], capture_output=True, text=True
    )
    filtered_result = subprocess.run(
        ["head", f"-{top_n}"], input=result.stdout, capture_output=True, text=True
    )

    # Sample output:
    # 43.642s systemd-networkd-wait-online.service
    lines = filtered_result.stdout.strip().splitlines()

    process_list = []
    for line in lines:
        process = re.search(r"\S+\s*$", line).group().strip()
        process_list.append({process: convert_to_secs(line)})

    boot_info = {"boot_time": boot_time, "longest_running_processes": process_list}

    return boot_info


def collect_resource_utilization():
    print("Collecting disk and memory usage...")

    # disk
    os_disk_usage = shutil.disk_usage("/")
    disk_usage = {
        "disk_size_gib": f"{os_disk_usage.total/1024**3:.2f}",
        "disk_usage_gib": f"{os_disk_usage.used/1024**3:.2f}",
    }

    # memory
    memory_info = psutil.virtual_memory()
    total_memory = memory_info.total // (1024**3)
    available_memory = memory_info.available // (1024**3)

    memory_usage = {
        "total_memory_gib": total_memory,
        "available_memory_gib": available_memory,
    }

    physical_cpu_count = psutil.cpu_count(logical=False)
    logical_cpu_count = psutil.cpu_count(logical=True)
    cpu_percent = psutil.cpu_percent()

    cpu_usage = {
        "physical_cpu_count": physical_cpu_count,
        "logical_cpu_count": logical_cpu_count,
        "cpu_percent": cpu_percent,
    }

    resource_utilization = {
        "disk_usage": disk_usage,
        "memory_usage": memory_usage,
        "cpu_usage": cpu_usage,
    }

    return resource_utilization


def collect_package_info():
    print("Collecting package info...")
    get_package_list = subprocess.run(
        ["rpm", "-qa"], capture_output=True, text=True, check=True
    )

    package_list = get_package_list.stdout.strip().splitlines()

    # TASK 4917: Adding package list resulted in hitting the size limit for the log,
    # so only logging package count until an alternative is implemented.
    package_info = {"package_count": len(package_list)}

    return package_info


def collect_cloud_init_info():
    print("Collecting cloud-init info...")

    # Collect cloud-init longest running processes
    result = subprocess.run(
        ["cloud-init", "analyze", "blame"], capture_output=True, text=True, check=True
    )

    lines = result.stdout.strip().splitlines()
    process_list = []
    top_n = 5

    # Skipping the first line as it is "-- Boot Record 01 --"
    # Skipping the last line as it is "x boot records analyzed"
    range = min(top_n + 1, len(lines) - 1)

    for line in lines[1:range]:
        record_details = line.split()
        if len(record_details) > 1:
            process_info = {}
            process_info["time"], process_info["process"] = record_details
            process_list.append(process_info)

    get_hostname = subprocess.run(
        ["hostname"], capture_output=True, text=True, check=True
    )

    cloud_init_info = {
        "hostname": get_hostname.stdout.strip(),
        "longest_running_processes": process_list,
    }

    return cloud_init_info


def get_selinux_mode():
    return subprocess.run(
        ["getenforce"], capture_output=True, text=True, check=True
    ).stdout.strip()


def collect_system_info():
    print("Collecting system info...")
    system_info = {"selinux_mode": get_selinux_mode()}
    return system_info


def get_asset_id():
    print("Collecting asset id...")
    
    return subprocess.run(
        ["cat", "/sys/devices/virtual/dmi/id/product_uuid"], capture_output=True, text=True
    ).stdout.lower().strip()


def has_valid_schema(data):
    schema_file = os.path.join(DATA_SCHEMA_DIR, DATA_SCHEMA_FILENAME)
    with open(schema_file, "r") as file:
        schema = json.load(file)

    try:
        jsonschema.validate(data, schema)
    except jsonschema.exceptions.ValidationError as err:
        print(f"Schema validation failed: {err}")
        return False
    return True


def main():
    print("Running azurelinux sysinfo collection...")
    asset_id = get_asset_id()
    os_info = collect_os_info()
    cloud_init_info = collect_cloud_init_info()
    boot_info = collect_boot_info()
    resource_utilization = collect_resource_utilization()
    package_info = collect_package_info()
    system_info = collect_system_info()

    # Use json as a data structure to store the data
    # since it is supported by Kusto
    data = {
        "$schema": f"{DATA_SCHEMA_VERSION}",
        "source": f"{SERVICE_NAME}",
        "asset_id": asset_id,
        "os_info": os_info,
        "cloud_init_info": cloud_init_info,
        "boot_info": boot_info,
        "resource_utilization": resource_utilization,
        "package_info": package_info,
        "system_info": system_info,
    }

    print(data)

    if has_valid_schema(data):
        # Dump the data to a log file, this path is added to fluentd config
        # and will be picked up by fluentd and sent through Geneva Agents
        with open(LOG_FILE_PATH, "w") as file:
            json.dump(data, file, separators=(',', ':'))

            # Add newline so that the fluentd tail plug-in consumes the log
            # line.
            file.write("\n")
        print("Azure Linux sysinfo collection completed successfully.")
    else:
        print("Azure Linux sysinfo collection failed.")
        exit(1)


if __name__ == "__main__":
    main()
